use std::sync::{
    Arc, 
    atomic::{AtomicUsize, Ordering},
};
use dashmap::DashMap;

#[derive(Debug)]
pub struct LuaMemoryTracker {
    limit: usize,
    per_worker: DashMap<u64, AtomicUsize>,
}

impl LuaMemoryTracker {
    pub fn new(limit: usize) -> Self {
        Self {
            limit,
            per_worker: DashMap::new(),
        }
    }

    pub fn get_limit(&self) -> usize {
        self.limit
    }

    pub fn get_memory_usage(&self, worker_id: u64) -> usize {
        self.per_worker
            .get(&worker_id)
            .map(|counter| counter.load(Ordering::Relaxed))
            .unwrap_or(0)
    }

    pub fn total_usage(&self) -> usize {
        self.per_worker
            .iter()
            .map(|entry| entry.load(Ordering::Relaxed))
            .sum()
    }

    pub fn update_usage(&self, worker_id: u64, usage: usize) {
        let entry = self.per_worker.entry(worker_id).or_insert(AtomicUsize::new(0));
        entry.store(usage, Ordering::Relaxed);
        
        // 检查是否超过限制
        if usage > self.limit {
            tracing::warn!(
                worker_id, 
                "Lua memory exceeded! Usage: {}/{} bytes", 
                usage, 
                self.limit
            );
        }
    }
}

impl mlua::allocator::Allocator for LuaMemoryTracker {
    fn alloc(&self, lua: &mlua::Lua, size: usize) -> mlua::Result<mlua::allocator::Alloc> {
        // 获取当前工作线程ID
        let worker_id = match lua.get_context_value::<u64>() {
            Ok(id) => id,
            Err(_) => 0,
        };

        // 计算新内存使用量
        let current = self.get_memory_usage(worker_id);
        let new_usage = current + size;

        // 检查是否超过限制
        if new_usage > self.limit {
            return Err(mlua::Error::AllocError {
                size,
                layout: None,
                extra: format!("Memory limit exceeded: {}/{}", new_usage, self.limit),
            });
        }

        // 更新使用量
        self.update_usage(worker_id, new_usage);

        // 实际分配内存
        mlua::allocator::System.alloc(lua, size)
    }

    fn dealloc(&self, lua: &mlua::Lua, size: usize, alloc: mlua::allocator::Alloc) {
        // 获取当前工作线程ID
        let worker_id = match lua.get_context_value::<u64>() {
            Ok(id) => id,
            Err(_) => 0,
        };

        // 更新使用量
        let current = self.get_memory_usage(worker_id);
        self.update_usage(worker_id, current.saturating_sub(size));

        // 实际释放内存
        mlua::allocator::System.dealloc(lua, size, alloc);
    }

    fn realloc(
        &self,
        lua: &mlua::Lua,
        old_size: usize,
        new_size: usize,
        alloc: mlua::allocator::Alloc,
    ) -> mlua::Result<mlua::allocator::Alloc> {
        // 获取当前工作线程ID
        let worker_id = match lua.get_context_value::<u64>() {
            Ok(id) => id,
            Err(_) => 0,
        };

        // 计算内存使用量变化
        let current = self.get_memory_usage(worker_id);
        let usage_change = new_size as isize - old_size as isize;
        let new_usage = if usage_change > 0 {
            current + usage_change as usize
        } else {
            current.saturating_sub(usage_change.abs() as usize)
        };

        // 检查是否超过限制
        if new_usage > self.limit {
            return Err(mlua::Error::AllocError {
                size: new_size,
                layout: None,
                extra: format!("Memory limit exceeded: {}/{}", new_usage, self.limit),
            });
        }

        // 更新使用量
        self.update_usage(worker_id, new_usage);

        // 实际重新分配内存
        mlua::allocator::System.realloc(lua, old_size, new_size, alloc)
    }
    // monitor.rs 线程安全保护
use std::sync::atomic::{AtomicUsize, Ordering};

struct LuaMemoryTracker {
    usage: DashMap<u64, AtomicUsize>,  // 替换原有 HashMap
    limit: usize
}

}